# 导入 FastAPI 相关模块，用于构建 API 和 WebSocket
from fastapi import FastAPI, WebSocket, HTTPException
# 导入 FastAPI 响应类型，用于流式响应和文件服务
from fastapi.responses import FileResponse
# 导入 CORS 中间件，支持跨域请求
from fastapi.middleware.cors import CORSMiddleware
# 导入静态文件服务模块
from fastapi.staticfiles import StaticFiles
# 导入 WebSocket 断开异常
from starlette.websockets import WebSocketDisconnect
# 导入系统操作模块，用于文件目录管理
import os
# 导入 Pydantic 模型，用于请求验证
from pydantic import BaseModel
# 导入异步事件循环模块
import asyncio
# 导入 JSON 处理模块
import json
# 导入 UUID 模块，生成唯一会话 ID
import uuid
# 导入类型注解模块
from typing import Optional
# 导入时间模块，记录处理时间
import time
# 导入正则表达式模块，用于匹配日常问候
import re
# 导入优化后的问答系统
from new_main import IntegratedQASystem


# 创建 FastAPI 应用实例，设置标题和描述
app = FastAPI(title="问答系统API", description="集成MySQL和RAG的智能问答系统")

# 配置 CORS 中间件，允许跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源（生产环境需限制）
    allow_credentials=True,  # 允许凭证
    allow_methods=["*"],  # 允许所有 HTTP 方法
    allow_headers=["*"],  # 允许所有头部
)

# 创建静态文件目录（如果不存在）
os.makedirs("static", exist_ok=True)

# 创建全局问答系统实例
qa_system = IntegratedQASystem()


# 定义日常问候用语模式和回复
GREETING_PATTERNS = [
    {
        "pattern": r"^(你好|您好|hi|hello)",  # 匹配问候语
        "response": "你好！我是黑马程序员，专注于为学生答疑解惑，很高兴为你服务！"
    },
    {
        "pattern": r"^(你是谁|您是谁|你叫什么|你的名字|who are you)",  # 匹配身份询问
        "response": "我是黑马程序员，你的智能学习助手，致力于提供 IT 教育相关的解答！"
    },
    {
        "pattern": r"^(在吗|在不在|有人吗)",  # 匹配在线确认
        "response": "我在！我是黑马程序员，随时为你解答问题！"
    },
    {
        "pattern": r"^(干嘛呢|你在干嘛|做什么)",  # 匹配状态询问
        "response": "我正在待命，随时为你解答 IT 学习相关的问题！有什么我可以帮你的？"
    }
]

# 定义查询请求模型
class QueryRequest(BaseModel):
    query: str  # 查询内容，必填
    source_filter: Optional[str] = None  # 学科过滤，可选
    session_id: Optional[str] = None  # 会话 ID，可选


# 挂载静态文件目录，服务前端文件
app.mount("/static", StaticFiles(directory="static"), name="static")

# 根路径重定向到 index.html
@app.get("/")
async def read_root():
    return FileResponse("static/index.html")


# 创建新会话接口
@app.post("/api/create_session")
async def create_session():
    session_id = str(uuid.uuid4())  # 生成唯一会话 ID
    return {"session_id": session_id}  # 返回会话 ID

# 查询历史消息接口
@app.get("/api/history/{session_id}")
async def get_history(session_id: str):
    try:
        # 获取指定会话的历史记录
        history = qa_system.get_session_history(session_id)
        # 返回会话 ID 和历史记录
        return {"session_id": session_id, "history": history}
    except Exception as e:
        # 抛出 HTTP 异常，包含错误信息
        raise HTTPException(status_code=500, detail=f"获取历史记录失败: {str(e)}")

# 清除历史消息接口
@app.delete("/api/history/{session_id}")
async def clear_history(session_id: str):
    # 清除指定会话的历史记录
    success = qa_system.clear_session_history(session_id)
    if success:
        # 返回成功状态，并包含 session_id 信息
        return {"status": "success", "message": "历史记录已清除", "session_id": session_id}
    else:
        # 抛出 HTTP 异常
        raise HTTPException(status_code=500, detail="清除历史记录失败")

# 检查是否为日常问候用语并返回模板回复
def check_greeting(query: str) -> Optional[str]:
    query_text = query.strip()  # 去除首尾空格
    for pattern_info in GREETING_PATTERNS:
        # 使用正则匹配，忽略大小写
        if re.match(pattern_info["pattern"], query_text, re.IGNORECASE):
            return pattern_info["response"]  # 返回匹配的回复
    return None  # 无匹配返回 None

# 非流式查询接口
@app.post("/api/query")
async def query(request: QueryRequest):
    start_time = time.time()  # 记录开始时间
    # 使用请求中的 session_id 或生成新 ID
    session_id = request.session_id or str(uuid.uuid4())
    # 检查是否为日常问候
    greeting_response = check_greeting(request.query)
    if greeting_response:
        # 返回问候回复
        return {
            "answer": greeting_response,
            "is_streaming": False,
            "session_id": session_id,
            "processing_time": time.time() - start_time
        }

    # 执行 BM25 搜索，安全处理返回值
    try:
        search_result = qa_system.bm25_search.search(request.query, threshold=0.85)
        # 处理解包问题 - 确保正确获取返回值
        if isinstance(search_result, tuple):
            if len(search_result) >= 2:
                answer, need_rag = search_result[0], search_result[1]
            elif len(search_result) == 1:
                answer, need_rag = search_result[0], False
            else:
                answer, need_rag = "未找到相关答案", False
        else:
            answer, need_rag = search_result, False
    except Exception as e:
        print(f"BM25搜索错误: {e}")
        answer, need_rag = "搜索过程中出现错误", False

    # 处理空答案情况，防止数据库插入空值错误
    if answer is None or str(answer).strip() == "":
        answer = "抱歉，我没有找到相关答案。请尝试重新表述您的问题或提供更多细节。"

    if need_rag:
        # 需要 RAG，提示使用 WebSocket
        return {
            "answer": "请使用WebSocket接口获取流式响应",
            "is_streaming": True,
            "session_id": session_id,
            "processing_time": time.time() - start_time
        }
    # 返回 MySQL 答案
    return {
        "answer": str(answer),
        "is_streaming": False,
        "session_id": session_id,
        "processing_time": time.time() - start_time
    }


# 流式查询 WebSocket 接口
@app.websocket("/api/stream")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()  # 接受 WebSocket 连接
    session_id = None
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            request_data = json.loads(data)  # 解析 JSON 数据
            # 获取查询参数
            query = request_data.get("query")
            source_filter = request_data.get("source_filter")
            session_id = request_data.get("session_id", str(uuid.uuid4()))

            # 验证会话有效性
            if not qa_system.is_valid_session(session_id):
                await websocket.send_json({
                    "type": "error",
                    "error": "无效的会话ID，请重新创建会话"
                })
                continue

            start_time = time.time()  # 记录开始时间

            # 发送开始标志
            await websocket.send_json({
                "type": "start",
                "session_id": session_id
            })

            # 检查是否为日常问候
            greeting_response = check_greeting(query)
            if greeting_response:
                # 发送问候回复
                await websocket.send_json({
                    "type": "token",
                    "token": greeting_response,
                    "session_id": session_id
                })
                # 发送结束标志
                await websocket.send_json({
                    "type": "end",
                    "session_id": session_id,
                    "is_complete": True,
                    "processing_time": time.time() - start_time
                })
                continue

            # 调用问答系统，流式处理查询
            collected_answer = ""
            try:
                # 使用同步 query 方法，并模拟流式输出
                result_generator = qa_system.query(query, source_filter=source_filter, session_id=session_id)

                for item in result_generator:
                    # 安全处理解包问题
                    if isinstance(item, tuple):
                        if len(item) >= 2:
                            token, is_complete = item[0], item[1]
                        elif len(item) == 1:
                            token, is_complete = item[0], False
                        else:
                            token, is_complete = "", True
                    else:
                        token, is_complete = item, False

                    # 确保 token 是字符串类型
                    token = str(token) if token is not None else ""
                    collected_answer += token  # 累积答案

                    if is_complete and not collected_answer.strip():
                        # 如果收集的答案为空，发送默认消息
                        await websocket.send_json({
                            "type": "token",
                            "token": "抱歉，我没有找到相关答案。请尝试重新表述您的问题或提供更多细节。",
                            "session_id": session_id
                        })
                        # 发送结束标志
                        await websocket.send_json({
                            "type": "end",
                            "session_id": session_id,
                            "is_complete": True,
                            "processing_time": time.time() - start_time
                        })
                        break
                    if token:
                        # 发送 token 数据
                        await websocket.send_json({
                            "type": "token",
                            "token": token,
                            "session_id": session_id
                        })
                    if is_complete:
                        # 发送结束标志
                        await websocket.send_json({
                            "type": "end",
                            "session_id": session_id,
                            "is_complete": True,
                            "processing_time": time.time() - start_time
                        })
                        break
                    await asyncio.sleep(0.01)  # 控制流式输出的速度
            except Exception as e:
                print(f"查询处理错误: {e}")
                # 发送错误消息
                await websocket.send_json({
                    "type": "error",
                    "error": f"处理查询时出现错误: {str(e)}",
                    "session_id": session_id
                })
                # 发送结束标志
                await websocket.send_json({
                    "type": "end",
                    "session_id": session_id,
                    "is_complete": True,
                    "processing_time": time.time() - start_time
                })
                break

    except WebSocketDisconnect as e:
        # 记录 WebSocket 断开信息
        print(f"WebSocket disconnected: code={e.code}, reason={e.reason}")
    except Exception as e:
        # 记录错误信息
        print(f"WebSocket error: {str(e)}")
        # 发送错误消息
        try:
            await websocket.send_json({
                "type": "error",
                "error": str(e),
                "session_id": session_id
            })
        except:
            pass
    finally:
        try:
            # 关闭 WebSocket 连接
            await websocket.close()
        except Exception as e:
            # 记录关闭连接时的错误
            print(f"Error closing WebSocket: {str(e)}")

# 健康检查接口
@app.get("/health")
async def health_check():
    return {"status": "healthy"}  # 返回健康状态


# 获取有效学科类别接口
@app.get("/api/sources")
async def get_sources():
    try:
        sources = qa_system.config.VALID_SOURCES
        return {"sources": sources}  # 返回学科类别列表
    except Exception as e:
        print(f"获取学科类别错误: {e}")
        return {"sources": []}  # 返回空列表作为默认值


# 主程序入口
if __name__ == "__main__":
    import uvicorn
    # 运行 FastAPI 应用，监听 0.0.0.0:8000
    uvicorn.run("app:app", host="127.0.0.1", port=8080, reload=False)
